import cv2
from PIL import Image
import numpy
import os

width  = 3840
height = 2160
widthDown1  = 1280
heightDown1 = 720
widthDown2  = 640
heightDown2 = 360
def mkdir(path): 
    folder = os.path.exists(path)
    if not folder:                   
        os.makedirs(path)            

def crop(frame):
    img = Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB))
    img_leftUp    = img.crop((0,0,int(width/2),int(height/2)))
    img_rightUp   = img.crop((int(width/2),0,width,int(height/2)))
    img_leftDown  = img.crop((0,int(height/2),int(width/2),height))
    img_rightDown = img.crop((int(width/2),int(height/2),width,height))
    img_leftUp    = img_leftUp.resize((widthDown1,heightDown1))
    img_rightUp   = img_rightUp.resize((widthDown1,heightDown1))
    img_leftDown  = img_leftDown.resize((widthDown1,heightDown1))
    img_rightDown = img_rightDown.resize((widthDown1,heightDown1))
    return img_leftUp,img_rightUp,img_leftDown,img_rightDown

def cropAndResize(frame):
    img = Image.fromarray(cv2.cvtColor(frame,cv2.COLOR_BGR2RGB))
    img_leftUp    = img.crop((0,0,int(width/2),int(height/2)))
    img_rightUp   = img.crop((int(width/2),0,width,int(height/2)))
    img_leftDown  = img.crop((0,int(height/2),int(width/2),height))
    img_rightDown = img.crop((int(width/2),int(height/2),width,height))
    img_leftUp    = img_leftUp.resize((widthDown2,heightDown2),resample=Image.NEAREST)
    img_rightUp   = img_rightUp.resize((widthDown2,heightDown2),resample=Image.NEAREST)
    img_leftDown  = img_leftDown.resize((widthDown2,heightDown2),resample=Image.NEAREST)
    img_rightDown = img_rightDown.resize((widthDown2,heightDown2),resample=Image.NEAREST)
    return img_leftUp,img_rightUp,img_leftDown,img_rightDown

NUM = 0
for filename in os.listdir(r'./source'):
    cap = cv2.VideoCapture('source/'+filename)  
    frameNum = cap.get(cv2.CAP_PROP_FRAME_COUNT)
    CENTER = int(frameNum / 2)

    cnt = 0
    while(cap.isOpened()):
        ret , frame = cap.read()
        if cnt >= CENTER - 23 and cnt <= CENTER - 17:
            if cnt == CENTER - 23:
                mkdir('data/'+str(NUM))
                mkdir('data/'+str(NUM+1))
                mkdir('data/'+str(NUM+2))
                mkdir('data/'+str(NUM+3))
            LU,RU,LD,RD = cropAndResize(frame)
            LU.save('data/'+str(NUM)+'/'+str(cnt-CENTER+23)+'.jpg')
            RU.save('data/'+str(NUM+1)+'/'+str(cnt-CENTER+23)+'.jpg')
            LD.save('data/'+str(NUM+2)+'/'+str(cnt-CENTER+23)+'.jpg')
            RD.save('data/'+str(NUM+3)+'/'+str(cnt-CENTER+23)+'.jpg')
            if cnt == CENTER - 20:
                ori_LU,ori_RU,ori_LD,ori_RD = crop(frame)
                ori_LU.save('label/'+str(NUM)+'.jpg')
                ori_RU.save('label/'+str(NUM+1)+'.jpg')
                ori_LD.save('label/'+str(NUM+2)+'.jpg')
                ori_RD.save('label/'+str(NUM+3)+'.jpg')
            if cnt == CENTER - 17:
                NUM += 4

        if cnt >= CENTER - 13 and cnt <= CENTER - 7:
            if cnt == CENTER - 13:
                mkdir('data/'+str(NUM))
                mkdir('data/'+str(NUM+1))
                mkdir('data/'+str(NUM+2))
                mkdir('data/'+str(NUM+3))
            LU,RU,LD,RD = cropAndResize(frame)
            LU.save('data/'+str(NUM)+'/'+str(cnt-CENTER+13)+'.jpg')
            RU.save('data/'+str(NUM+1)+'/'+str(cnt-CENTER+13)+'.jpg')
            LD.save('data/'+str(NUM+2)+'/'+str(cnt-CENTER+13)+'.jpg')
            RD.save('data/'+str(NUM+3)+'/'+str(cnt-CENTER+13)+'.jpg')
            if cnt == CENTER - 10:
                ori_LU,ori_RU,ori_LD,ori_RD = crop(frame)
                ori_LU.save('label/'+str(NUM)+'.jpg')
                ori_RU.save('label/'+str(NUM+1)+'.jpg')
                ori_LD.save('label/'+str(NUM+2)+'.jpg')
                ori_RD.save('label/'+str(NUM+3)+'.jpg')
            if cnt == CENTER - 7:
                NUM += 4

        if cnt >= CENTER - 3 and cnt <= CENTER + 3:
            if cnt == CENTER - 3:
                mkdir('data/'+str(NUM))
                mkdir('data/'+str(NUM+1))
                mkdir('data/'+str(NUM+2))
                mkdir('data/'+str(NUM+3))
            LU,RU,LD,RD = cropAndResize(frame)
            LU.save('data/'+str(NUM)+'/'+str(cnt-CENTER+3)+'.jpg')
            RU.save('data/'+str(NUM+1)+'/'+str(cnt-CENTER+3)+'.jpg')
            LD.save('data/'+str(NUM+2)+'/'+str(cnt-CENTER+3)+'.jpg')
            RD.save('data/'+str(NUM+3)+'/'+str(cnt-CENTER+3)+'.jpg')
            if cnt == CENTER:
                ori_LU,ori_RU,ori_LD,ori_RD = crop(frame)
                ori_LU.save('label/'+str(NUM)+'.jpg')
                ori_RU.save('label/'+str(NUM+1)+'.jpg')
                ori_LD.save('label/'+str(NUM+2)+'.jpg')
                ori_RD.save('label/'+str(NUM+3)+'.jpg')
            if cnt == CENTER + 3:
                NUM += 4

        if cnt >= CENTER + 7 and cnt <= CENTER + 13:
            if cnt == CENTER + 7:
                mkdir('data/'+str(NUM))
                mkdir('data/'+str(NUM+1))
                mkdir('data/'+str(NUM+2))
                mkdir('data/'+str(NUM+3))
            LU,RU,LD,RD = cropAndResize(frame)
            LU.save('data/'+str(NUM)+'/'+str(cnt-CENTER-7)+'.jpg')
            RU.save('data/'+str(NUM+1)+'/'+str(cnt-CENTER-7)+'.jpg')
            LD.save('data/'+str(NUM+2)+'/'+str(cnt-CENTER-7)+'.jpg')
            RD.save('data/'+str(NUM+3)+'/'+str(cnt-CENTER-7)+'.jpg')
            if cnt == CENTER + 10:
                ori_LU,ori_RU,ori_LD,ori_RD = crop(frame)
                ori_LU.save('label/'+str(NUM)+'.jpg')
                ori_RU.save('label/'+str(NUM+1)+'.jpg')
                ori_LD.save('label/'+str(NUM+2)+'.jpg')
                ori_RD.save('label/'+str(NUM+3)+'.jpg')
            if cnt == CENTER + 13:
                NUM += 4

        if cnt >= CENTER + 17 and cnt <= CENTER + 23:
            if cnt == CENTER + 17:
                mkdir('data/'+str(NUM))
                mkdir('data/'+str(NUM+1))
                mkdir('data/'+str(NUM+2))
                mkdir('data/'+str(NUM+3))
            LU,RU,LD,RD = cropAndResize(frame)
            LU.save('data/'+str(NUM)+'/'+str(cnt-CENTER-17)+'.jpg')
            RU.save('data/'+str(NUM+1)+'/'+str(cnt-CENTER-17)+'.jpg')
            LD.save('data/'+str(NUM+2)+'/'+str(cnt-CENTER-17)+'.jpg')
            RD.save('data/'+str(NUM+3)+'/'+str(cnt-CENTER-17)+'.jpg')
            if cnt == CENTER + 20:
                ori_LU,ori_RU,ori_LD,ori_RD = crop(frame)
                ori_LU.save('label/'+str(NUM)+'.jpg')
                ori_RU.save('label/'+str(NUM+1)+'.jpg')
                ori_LD.save('label/'+str(NUM+2)+'.jpg')
                ori_RD.save('label/'+str(NUM+3)+'.jpg')
            if cnt == CENTER + 23:
                NUM += 4

        cnt += 1
        if ret == False:
            break
    cap.release()
    # cv2.destroyAllWindows()